import re
import time

import requests
import pandas as pd
from tqdm import tqdm


def extract_uniprot_id(text):
    """从混合字符串中提取标准UniProt ID（如B4FS35）"""
    match = re.search(r'([A-N,R-Z][0-9][A-Z][A-Z,0-9][A-Z,0-9][0-9])', text)
    return match.group(1) if match else None


def fetch_sequences(input_excel, output_file='combined.fasta'):
    """
    从Excel提取混合格式ID中的UniProt ID并下载序列

    参数:
        input_excel: 包含混合ID的Excel文件路径
        output_file: 输出FASTA文件名
    """
    # 读取Excel并提取有效ID
    try:
        df = pd.read_excel(input_excel, header=None)
        raw_ids = df.iloc[:, 0].dropna().tolist()
        uniprot_ids = [extract_uniprot_id(s) for s in raw_ids]
        uniprot_ids = list(filter(None, uniprot_ids))  # 过滤无效ID
    except Exception as e:
        print(f"文件处理错误: {str(e)}")
        return

    # 下载序列
    success_count = 0
    failed_ids = []

    with open(output_file, 'w') as f_out:
        for uid in tqdm(uniprot_ids, desc="下载进度"):
            try:
                url = f"https://www.uniprot.org/uniprot/{uid}.fasta"
                response = requests.get(url, timeout=10)
                time.sleep(0.34)  # 遵守3次/秒的API限制

                if response.status_code == 200:
                    f_out.write(response.text)
                    success_count += 1
                else:
                    failed_ids.append(uid)
            except Exception as e:
                print(f"ID {uid} 下载失败: {str(e)}")
                failed_ids.append(uid)

    # 生成报告
    print(f"\n成功下载 {success_count} 条序列到 {output_file}")
    if failed_ids:
        with open('failed_ids.txt', 'w') as f:
            f.write('\n'.join(failed_ids))
        print(f"失败ID已保存到 failed_ids.txt")


if __name__ == "__main__":
    fetch_sequences("mixed_ids.xlsx")

